from robot_arm_module import RobotArmModule
import numpy as np

def find_max_torques():
    # 创建机械臂实例 L1=114, L2=60, 末端负载1g
    arm = RobotArmModule(l1=114, l2=60, mass=0.001)
    
    # 定义工作空间范围和步长
    x_min, x_max = 50, 150
    y_min, y_max = 0, 50
    step = 1  # 步长，越小精度越高但计算时间越长
    
    # 初始化最大值跟踪变量
    max_torque1 = -float('inf')
    max_torque2 = -float('inf')
    
    # 存储最大值对应的状态
    max_t1_info = None
    max_t2_info = None
    
    # 计算总点数和初始化进度计数器
    x_points = int((x_max - x_min) / step) + 1
    y_points = int((y_max - y_min) / step) + 1
    total_points = x_points * y_points
    count = 0
    
    print(f"开始遍历工作空间，共 {total_points} 个点...")
    
    # 遍历工作空间内的所有点
    for x in np.arange(x_min, x_max + step, step):
        for y in np.arange(y_min, y_max + step, step):
            # 计算逆运动学（肘上解）
            angles = arm.inverse_kinematics(x, y, solution_type=0)
            
            # 检查是否有解
            if angles:
                theta1_ik, theta2_ik = angles
                
                # 计算力矩
                torque1, torque2 = arm.calculate_torques(theta1_ik, theta2_ik)
                
                # 取力矩绝对值（因为力矩方向可能为负，但我们关心大小）
                abs_torque1 = abs(torque1)
                abs_torque2 = abs(torque2)
                
                # 检查是否为新的最大值
                if abs_torque1 > max_torque1:
                    max_torque1 = abs_torque1
                    max_t1_info = {
                        'x': x,
                        'y': y,
                        'theta1': theta1_ik,
                        'theta2': theta2_ik,
                        'torque': torque1
                    }
                
                if abs_torque2 > max_torque2:
                    max_torque2 = abs_torque2
                    max_t2_info = {
                        'x': x,
                        'y': y,
                        'theta1': theta1_ik,
                        'theta2': theta2_ik,
                        'torque': torque2
                    }
            
            # 显示进度
            count += 1
            if count % 1000 == 0:
                progress = (count / total_points) * 100
                print(f"进度: {progress:.1f}% ({count}/{total_points})")
    
    print("\n遍历完成！")
    print("\n=== 关节1最大力矩信息 ===")
    if max_t1_info:
        print(f"位置: ({max_t1_info['x']}, {max_t1_info['y']})")
        print(f"关节角度: theta1={max_t1_info['theta1']:.2f}°, theta2={max_t1_info['theta2']:.2f}°")
        print(f"最大力矩值: {max_t1_info['torque']:.6f} N·m (绝对值: {max_torque1:.6f} N·m)")
    
    print("\n=== 关节2最大力矩信息 ===")
    if max_t2_info:
        print(f"位置: ({max_t2_info['x']}, {max_t2_info['y']})")
        print(f"关节角度: theta1={max_t2_info['theta1']:.2f}°, theta2={max_t2_info['theta2']:.2f}°")
        print(f"最大力矩值: {max_t2_info['torque']:.6f} N·m (绝对值: {max_torque2:.6f} N·m)")

if __name__ == "__main__":
    find_max_torques()
